#define _GNU_SOURCE 1
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ipc.h>
#include <sys/mman.h>
#include <sys/sem.h>
#include <sys/shm.h>
#include <sys/types.h>
#include <time.h>

#include "hashpipe_databuf.h"
#include "hashpipe_error.h"
#include "hashpipe_ipckey.h"

#ifndef SHM_HUGE_SHIFT
#define SHM_HUGE_SHIFT 26
#endif
#ifndef SHM_HUGE_2MB
#define SHM_HUGE_2MB (21 << SHM_HUGE_SHIFT)
#endif
#ifndef SHM_HUGE_1GB
#define SHM_HUGE_1GB (30 << SHM_HUGE_SHIFT)
#endif

/* union for semaphore ops. */
union semun {
    int val;
    struct semid_ds* buf;
    unsigned short* array;
    struct seminfo* __buf;
};

/*
(1) shmget: create or get shared memory.
(2) shmat: shared memory <---> local memory.
(3) shmctl: lock shared memory, prevents swapping out, improves performance.
(4) semget: create or obtain a set of semaphores.
(5) semctl: modify semaphores, set it be zero.
*/

hashpipe_databuf_t* hashpipe_databuf_create(int instance_id,
    int databuf_id, size_t header_size, size_t block_size, int n_block)
{
    int rv = 0;
    int verify_sizing = 0;
    int newly_created = 0;

    size_t total_size = header_size + block_size * n_block;
    size_t total_size_page_aligned;

    if (header_size < sizeof(hashpipe_databuf_t)) {
        hashpipe_error(__FUNCTION__, "header size must be larger than %lu",
            sizeof(hashpipe_databuf_t));
        return NULL;
    }

    /* Get shared memory block */
    key_t key = hashpipe_databuf_key(instance_id);
    if (key == HASHPIPE_KEY_ERROR) {
        hashpipe_error(__FUNCTION__, "hashpipe_databuf_key error");
        return NULL;
    }

    int shmid;
    // First try 1 GB pages
    hashpipe_info(__FUNCTION__, "total_size %lu", total_size);
    total_size_page_aligned = total_size + ((-total_size) % (1 << 30));
    hashpipe_info(__FUNCTION__, "total_size 1GB aligned %lu (%lx)",
        total_size_page_aligned, total_size_page_aligned);

    shmid = shmget(key + databuf_id - 1, total_size_page_aligned,
        0666 | IPC_CREAT | IPC_EXCL | SHM_HUGETLB | SHM_HUGE_1GB);

    if (shmid != -1) {
        hashpipe_info(__FUNCTION__, "created shared memory for key %08x with 1GB pages",
            key + databuf_id - 1);
        newly_created = 1;
    } else if (errno == EEXIST) {
        hashpipe_info(__FUNCTION__, "shared memory key %08x already exists",
            key + databuf_id - 1);
        // Already exists, call shmget again without IPC_CREAT
        shmid = shmget(key + databuf_id - 1, total_size, 0666);
        // Verify buffer sizing
        verify_sizing = 1;
    } else if (errno == ENOMEM) {
        // Try with 2MB pages
        total_size_page_aligned = total_size + ((-total_size) % (1 << 21));
        hashpipe_info(__FUNCTION__, "total_size 2MB aligned %lu (%lx)",
            total_size_page_aligned, total_size_page_aligned);
        shmid = shmget(key + databuf_id - 1, total_size_page_aligned,
            0666 | IPC_CREAT | IPC_EXCL | SHM_HUGETLB | SHM_HUGE_2MB);

        if (shmid != -1) {
            hashpipe_info(__FUNCTION__,
                "created shared memory for key %08x with 2MB pages",
                key + databuf_id - 1);
            newly_created = 1;
        } else if (errno == ENOMEM) {
            // Try without huge pages
            shmid = shmget(key + databuf_id - 1, total_size,
                0666 | IPC_CREAT | IPC_EXCL);
            if (shmid != -1) {
                hashpipe_info(__FUNCTION__,
                    "created shared memory for key %08x without huge pages",
                    key + databuf_id - 1);
                newly_created = 1;
            }
        }
    }
    if (shmid == -1) {
        perror("shmget");
        hashpipe_error(__FUNCTION__, "shmget error");
        return NULL;
    }

    // attach
    hashpipe_databuf_t* d;
    d = shmat(shmid, NULL, 0);
    if (d == (void*)-1) {
        hashpipe_error(__FUNCTION__, "shmat error");
        return NULL;
    }

    if (verify_sizing) {
        // Make sure existing sizes match expectaions
        if (d->header_size != header_size
            || d->block_size != block_size
            || d->n_block != n_block) {
            char msg[256];
            sprintf(msg, "existing databuf size mismatch "
                         "(%lu + %lu x %d) != (%lu + %ld x %d)",
                d->header_size, d->block_size, d->n_block,
                header_size, block_size, n_block);
            hashpipe_error(__FUNCTION__, msg);
            if (shmdt(d)) {
                hashpipe_error(__FUNCTION__, "shmdt error");
            }
            return NULL;
        }
    }

    // try to lock in memory
    rv = shmctl(shmid, SHM_LOCK, NULL);
    if (rv == -1) {
        perror("shmctl");
        hashpipe_error(__FUNCTION__, "Error locking shared memory.");
        return NULL;
    }

    if (newly_created) {
        /* Zero out newly created databuf */
        memset(d, 0, total_size);

        /* Fill params into databuf */
        d->shmid = shmid;
        d->semid = 0;
        d->header_size = header_size;
        d->n_block = n_block;
        d->block_size = block_size;
        sprintf(d->data_type, "unknown");
    }

    /* Get semaphores set up */
    d->semid = semget(key + databuf_id - 1, n_block, 0666 | IPC_CREAT);
    if (d->semid == -1) {
        hashpipe_error(__FUNCTION__, "semget error");
        return NULL;
    }

    /* Init semaphores to 0 */
    union semun arg;
    arg.array = (unsigned short*)malloc(sizeof(unsigned short) * n_block);
    memset(arg.array, 0, sizeof(unsigned short) * n_block);
    rv = semctl(d->semid, 0, SETALL, arg);
    if (rv == -1) {
        perror("semctl");
        hashpipe_error(__FUNCTION__, "Error clearing semaphores.");
        free(arg.array);
        return NULL;
    }
    free(arg.array);

    return d;
}

int hashpipe_databuf_detach(hashpipe_databuf_t* d)
{
    if (d) {
        int rv = shmdt(d);
        if (rv != 0) {
            hashpipe_error(__FUNCTION__, "shmdt error");
            return HASHPIPE_ERR_SYS;
        }
    }
    return HASHPIPE_OK;
}

void hashpipe_databuf_clear(hashpipe_databuf_t* d)
{
    /* Zero out semaphores */
    union semun arg;
    arg.array = (unsigned short*)malloc(sizeof(unsigned short) * d->n_block);
    memset(arg.array, 0, sizeof(unsigned short) * d->n_block);
    semctl(d->semid, 0, SETALL, arg);
    free(arg.array);
}

char* hashpipe_databuf_data(hashpipe_databuf_t* d, int block_id)
{
    return (char*)d + d->header_size + d->block_size * block_id;
}

hashpipe_databuf_t* hashpipe_databuf_attach(int instance_id, int databuf_id)
{
    /* Get shmid */
    key_t key = hashpipe_databuf_key(instance_id);
    if (key == HASHPIPE_KEY_ERROR) {
        hashpipe_error(__FUNCTION__, "hashpipe_databuf_key error");
        return NULL;
    }
    int shmid;
    shmid = shmget(key + databuf_id - 1, 0, 0666);
    if (shmid == -1) {
        // Doesn't exist, exit quietly otherwise complain
        if (errno != ENOENT)
            hashpipe_error(__FUNCTION__, "shmget error");
        return NULL;
    }

    /* Attach */
    hashpipe_databuf_t* d;
    d = shmat(shmid, NULL, 0);
    if (d == (void*)-1) {
        hashpipe_error(__FUNCTION__, "shmat error");
        return NULL;
    }

    return d;
}

int hashpipe_databuf_total_status(hashpipe_databuf_t* d)
{
    /* Get all values at once */
    union semun arg;
    arg.array = (unsigned short*)malloc(sizeof(unsigned short) * d->n_block);
    memset(arg.array, 0, sizeof(unsigned short) * d->n_block);
    semctl(d->semid, 0, GETALL, arg);
    int tot = 0;
    for (int i = 0; i < d->n_block; i++)
        tot += arg.array[i];
    free(arg.array);
    return tot;
}

uint64_t hashpipe_databuf_total_mask(hashpipe_databuf_t* d)
{
    /* Get all values at once */
    union semun arg;
    arg.array = (unsigned short*)malloc(sizeof(unsigned short) * d->n_block);
    memset(arg.array, 0, sizeof(unsigned short) * d->n_block);
    semctl(d->semid, 0, GETALL, arg);

    int n = d->n_block;
    if (n > 64)
        n = 64;
    uint64_t tot = 0;
    for (int i = 0; i < n; i++) {
        if (arg.array[i]) {
            tot |= (1 << i);
        }
    }
    free(arg.array);
    return tot;
}

int hashpipe_databuf_wait_free_timeout(hashpipe_databuf_t* d,
    int block_id, struct timespec* timeout)
{
    int rv;
    struct sembuf op;
    op.sem_num = block_id;
    op.sem_op = 0; // wait for the signal size to change to 0.
    op.sem_flg = 0;
    rv = semtimedop(d->semid, &op, 1, timeout);
    if (rv == -1) {
        if (errno == EAGAIN) {
            printf("%s(%p, %d) timeout (%016lx)\n",
                __FUNCTION__, d, block_id, hashpipe_databuf_total_mask(d));
            return HASHPIPE_TIMEOUT;
        }
        // Don't complain on a signal interruption
        if (errno == EINTR)
            return HASHPIPE_ERR_SYS;
        hashpipe_error(__FUNCTION__, "semop error");
        perror("semop");
        return HASHPIPE_ERR_SYS;
    }
    return 0;
}

int hashpipe_databuf_wait_free(hashpipe_databuf_t* d, int block_id)
{
    struct timespec timeout;
    timeout.tv_sec = 0;
    timeout.tv_nsec = 250000000;
    return hashpipe_databuf_wait_free_timeout(d, block_id, &timeout);
}

int hashpipe_databuf_busywait_free(hashpipe_databuf_t* d, int block_id)
{
    int rv;
    struct sembuf op;
    op.sem_num = block_id;
    op.sem_op = 0;
    op.sem_flg = IPC_NOWAIT;

    do {
        rv = semop(d->semid, &op, 1);
    } while (rv == -1 && errno == EAGAIN);

    if (rv == -1) {
        // Don't complain on a signal interruption
        if (errno == EINTR)
            return HASHPIPE_ERR_SYS;
        hashpipe_error(__FUNCTION__, "semop error");
        perror("semop");
        return HASHPIPE_ERR_SYS;
    }
    return 0;
}

int hashpipe_databuf_wait_filled_timeout(hashpipe_databuf_t* d,
    int block_id, struct timespec* timeout)
{
    /* This needs to wait for the semval of the given block
     * to become > 0, but NOT immediately decrement it to 0.
     * Probably do this by giving an array of semops, since
     * (afaik) the whole array happens atomically:
     * step 1: wait for val=1 then decrement (semop=-1)
     * step 2: increment by 1 (semop=1)
     */
    int rv;
    struct sembuf op[2];
    op[0].sem_num = op[1].sem_num = block_id;
    op[0].sem_flg = op[1].sem_flg = 0;
    op[0].sem_op = -1;
    op[1].sem_op = 1;
    rv = semtimedop(d->semid, op, 2, timeout);
    if (rv == -1) {
        if (errno == EAGAIN)
            return HASHPIPE_TIMEOUT;
        // Don't complain on a signal interruption
        if (errno == EINTR)
            return HASHPIPE_ERR_SYS;
        hashpipe_error(__FUNCTION__, "semop error");
        perror("semop");
        return HASHPIPE_ERR_SYS;
    }
    return 0;
}

int hashpipe_databuf_wait_filled(hashpipe_databuf_t* d, int block_id)
{
    struct timespec timeout;
    timeout.tv_sec = 0;
    timeout.tv_nsec = 250000000;

    return hashpipe_databuf_wait_filled_timeout(d, block_id, &timeout);
}

int hashpipe_databuf_busywait_filled(hashpipe_databuf_t* d, int block_id)
{
    /* This needs to wait for the semval of the given block
     * to become > 0, but NOT immediately decrement it to 0.
     * Probably do this by giving an array of semops, since
     * (afaik) the whole array happens atomically:
     * step 1: wait for val=1 then decrement (semop=-1)
     * step 2: increment by 1 (semop=1)
     */
    int rv;
    struct sembuf op[2];
    op[0].sem_num = op[1].sem_num = block_id;
    op[0].sem_flg = IPC_NOWAIT;
    op[1].sem_flg = IPC_NOWAIT;
    op[0].sem_op = -1;
    op[1].sem_op = 1;

    do {
        rv = semop(d->semid, op, 2);
    } while (rv == -1 && errno == EAGAIN);

    if (rv == -1) {
        // Don't complain on a signal interruption
        if (errno == EINTR)
            return HASHPIPE_ERR_SYS;
        hashpipe_error(__FUNCTION__, "semop error");
        perror("semop");
        return HASHPIPE_ERR_SYS;
    }
    return 0;
}

int hashpipe_databuf_set_free(hashpipe_databuf_t* d, int block_id)
{
    /* This function should always succeed regardless of the current
     * state of the specified databuf.  So we use semctl (not semop) to set
     * the value to zero.
     */
    int rv;
    union semun arg;
    arg.val = 0;
    rv = semctl(d->semid, block_id, SETVAL, arg);
#ifdef HASHPIPE_TRACE
    printf("after %s(%p, %d) %016lx\n",
        __FUNCTION__, d, block_id, hashpipe_databuf_total_mask(d));
#endif
    if (rv == -1) {
        hashpipe_error(__FUNCTION__, "semctl error");
        return HASHPIPE_ERR_SYS;
    }
    return 0;
}

int hashpipe_databuf_set_filled(hashpipe_databuf_t* d, int block_id)
{
    /* This function should always succeed regardless of the current
     * state of the specified databuf.  So we use semctl (not semop) to set
     * the value to one.
     */
    int rv;
    union semun arg;
    arg.val = 1;
    rv = semctl(d->semid, block_id, SETVAL, arg);
#ifdef HASHPIPE_TRACE
    printf("after %s(%p, %d) %016lx\n",
        __FUNCTION__, d, block_id, hashpipe_databuf_total_mask(d));
#endif
    if (rv == -1) {
        hashpipe_error(__FUNCTION__, "semctl error");
        return HASHPIPE_ERR_SYS;
    }
    return 0;
}
